import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, RandomizedSearchCV, StratifiedKFold
from sklearn.preprocessing import LabelEncoder, StandardScaler
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score, roc_auc_score, classification_report
import matplotlib.pyplot as plt
import seaborn as sns

# ================ 数据加载与预处理 ================
try:
    # 读取训练数据
    data = pd.read_csv('../../data/raw/train.csv')
    print("数据加载成功，形状：", data.shape)
except FileNotFoundError:
    print("文件未找到，请检查路径是否正确！")
    exit()

# 修正特征列名（处理可能的拼写错误）
feature_names = {'Age': 'Age',  # 年龄
    'BusinessTravel': 'BusinessTravel',  # 商务旅行情况
    'Department': 'Department',  # 部门
    'DistanceFromHome': 'DistanceFromHome',  # 家与工作地点的距离
    'Education': 'Education',  # 教育程度
    'EducationField': 'EducationField',  # 教育领域
    'EnvironmentSatisfaction': 'EnvironmentSatisfaction',  # 工作环境满意度
    'Gender': 'Gender',  # 性别
    'JobInvolvement': 'JobInvolvement',  # 工作投入度
    'JobLevel': 'JobLevel',  # 工作级别
    'JobRole': 'JobRole',  # 工作角色
    'JobSatisfaction': 'JobSatisfaction',  # 工作满意度
    'MaritalStatus': 'MaritalStatus',  # 婚姻状况
    'MonthlyIncome': 'MonthlyIncome',  # 月收入
    'NumCompaniesWorked': 'NumCompaniesWorked',  # 曾工作过的公司数量
    'OverTime': 'OverTime',  # 是否加班
    'PercentSalaryHike': 'PercentSalaryHike',  # 薪资涨幅百分比
    'PerformanceRating': 'PerformanceRating',  # 绩效评级
    'RelationshipSatisfaction': 'RelationshipSatisfaction',  # 人际关系满意度
    'StockOptionLevel': 'StockOptionLevel',  # 股票期权水平
    'TotalWorkingYears': 'TotalWorkingYears',  # 总工作年限
    'TrainingTimesLastYear': 'TrainingTimesLastYear',  # 去年参加培训次数
    'WorkLifeBalance': 'WorkLifeBalance',  # 工作与生活平衡度
    'YearsAtCompany': 'YearsAtCompany',  # 在公司工作年限
    'YearsInCurrentRole': 'YearsInCurrentRole',  # 在当前岗位工作年限
    'YearsSinceLastPromotion': 'YearsSinceLastPromotion',  # 自上次晋升以来的年限
    'YearsWithCurrManager': 'YearsWithCurrManager'  # 与当前经理共事年限
}

# 验证并选择有效的特征列
actual_columns = data.columns.tolist()
valid_features = []

for key, value in feature_names.items():
    if value in actual_columns:
        valid_features.append(value)
    elif key in actual_columns:
        valid_features.append(key)
        print(f"警告：发现列名 '{key}'，可能是拼写错误，已使用。建议使用标准名称：'{value}'")
    else:
        print(f"错误：列名 '{value}' 不存在于数据中！")

if not valid_features:
    print("没有找到有效的特征列！")
    exit()

print("使用的特征列：", valid_features)

# 提取特征和目标变量
x = data[valid_features]
y = data['Attrition']  # 离职率（目标变量）

# 数据预处理
# 1. 检查并处理缺失值
missing_values = x.isnull().sum()
if missing_values.sum() > 0:
    print("\n缺失值统计：")
    print(missing_values[missing_values > 0])
    # 对数值型特征用均值填充缺失值
    x = x.fillna(x.mean())
else:
    print("\n数据无缺失值")

# 2. 对分类变量进行编码（如果有）
categorical_cols = x.select_dtypes(include=['object']).columns
if not categorical_cols.empty:
    print("\n需要编码的分类变量：", categorical_cols.tolist())
    for col in categorical_cols:
        # 使用标签编码将分类变量转换为数值
        le = LabelEncoder()
        x[col] = le.fit_transform(x[col])
else:
    print("\n数据中没有分类变量需要编码")

# ================ 数据划分与标准化 ================
# 划分训练集和测试集（80%训练，20%测试）
# 使用stratify参数确保训练集和测试集的类别分布一致
x_train, x_test, y_train, y_test = train_test_split(
    x, y, test_size=0.2, random_state=25, stratify=y
)

# 标准化数值特征（使特征具有零均值和单位方差）
# 注意：XGBoost对特征缩放不敏感，但标准化有助于某些评估指标
scaler = StandardScaler()
x_train_scaled = scaler.fit_transform(x_train)
x_test_scaled = scaler.transform(x_test)

# ================ 模型定义与超参数优化 ================
# 定义XGBoost分类器
# objective='binary:logistic' 适用于二分类问题
# eval_metric='auc' 使用AUC作为评估指标
xgb = XGBClassifier(
    objective='binary:logistic',
    eval_metric='auc',
    random_state=25,
    use_label_encoder=False  # 避免警告
)

# 定义超参数搜索空间
# 使用对数分布和合理范围来提高搜索效率
param_dist = {
    'learning_rate': np.logspace(-2,0,10),  # 学习率：0.01到1
    'max_depth': list(range(3, 10)),          # 树的最大深度
    'n_estimators': list(range(100,300,50)),  # 树的数量
    'subsample': np.linspace(0.6, 1.0, 5),    # 样本采样比例
    'colsample_bytree': np.linspace(0.6, 1.0, 5),  # 特征采样比例
    'reg_alpha': np.logspace(-3, 1, 5),       # L1正则化（减少过拟合）
    'reg_lambda': np.logspace(-3, 1, 5),      # L2正则化（减少过拟合）
    'min_child_weight': [1, 3, 5, 7]          # 子节点所需的最小样本权重和
}

# 创建分层K折交叉验证
# StratifiedKFold确保每一折中类别分布与原始数据一致
stratified_kfold = StratifiedKFold(n_splits=5, shuffle=True, random_state=25)

# 使用随机搜索进行超参数优化
# 相比网格搜索，随机搜索效率更高，尤其是参数空间较大时
random_search = RandomizedSearchCV(
    estimator=xgb,
    param_distributions=param_dist,
    n_iter=50,             # 尝试50种不同的参数组合
    scoring='roc_auc',     # 使用ROC AUC作为评分指标
    cv=stratified_kfold,   # 使用分层K折交叉验证
    n_jobs=-1,             # 使用所有CPU核心并行计算
    verbose=2,             # 显示详细的搜索进度
    random_state=25
)

print("\n开始随机搜索最佳参数...")
random_search.fit(x_train_scaled, y_train)

# ================ 模型评估与结果分析 ================
# 输出最佳参数和性能
print("\n最佳参数:")
for param, value in random_search.best_params_.items():
    print(f"  {param}: {value}")

print(f"最佳交叉验证 ROC AUC: {random_search.best_score_:.4f}")

# 使用最佳模型进行预测
best_model = random_search.best_estimator_
y_pred = best_model.predict(x_test_scaled)
y_pred_proba = best_model.predict_proba(x_test_scaled)[:, 1]  # 获取正类概率

# 评估模型性能
print(f"\n测试集准确率: {accuracy_score(y_test, y_pred):.4f}")
print(f"测试集 ROC AUC: {roc_auc_score(y_test, y_pred_proba):.4f}")

print("\n分类报告:")
print(classification_report(y_test, y_pred))

# 特征重要性分析
importance_df = pd.DataFrame({
    'Feature': x.columns,
    'Importance': best_model.feature_importances_
}).sort_values('Importance', ascending=False)

print("\n特征重要性:")
print(importance_df)

# 可视化特征重要性
plt.figure(figsize=(10, 6))
sns.barplot(x='Importance', y='Feature', data=importance_df)
plt.title('XGBoost 特征重要性')
plt.tight_layout()
plt.savefig('feature_importance.png')
plt.show()

# 学习曲线分析（如果模型支持记录训练过程）
results = best_model.evals_result()
if 'validation_0' in results:
    plt.figure(figsize=(10, 6))
    plt.plot(results['validation_0']['auc'], label='Train')
    plt.plot(results['validation_1']['auc'], label='Test')
    plt.xlabel('Boosting Iterations')
    plt.ylabel('AUC')
    plt.title('XGBoost 学习曲线')
    plt.legend()
    plt.savefig('learning_curve.png')
    plt.show()